import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset

# 设置随机种子确保结果可复现
torch.manual_seed(42)
np.random.seed(42)

# 定义GRU模型
class GRUModel(nn.Module):
    def __init__(self, input_dim, hidden_dim, output_dim, num_layers=1, dropout=0.0):
        super(GRUModel, self).__init__()
        self.hidden_dim = hidden_dim
        self.num_layers = num_layers
        
        # GRU层
        self.gru = nn.GRU(
            input_dim, 
            hidden_dim, 
            num_layers, 
            batch_first=True,  # 输入格式为(batch, seq_len, feature)
            dropout=dropout if num_layers > 1 else 0
        )
        
        # 全连接层，将GRU的输出映射到预测维度
        self.fc = nn.Linear(hidden_dim, output_dim)
        
    def forward(self, x):
        # 初始化隐藏状态
        h0 = torch.zeros(self.num_layers, x.size(0), self.hidden_dim).to(x.device)
        
        # GRU前向传播
        out, _ = self.gru(x, h0)
        
        # 我们只取序列中的最后一个时间步的输出
        # 假设out的形状是(batch_size, seq_len, hidden_dim)
        out = self.fc(out[:, -1, :])  # (32,1)
        return out

# 生成示例数据（时间序列预测任务）
def generate_data(n_samples=1000, seq_len=10, input_dim=1):
    """生成简单的时间序列数据，其中目标值是输入序列的总和"""
    # 生成随机序列作为输入
    X = np.random.randn(n_samples, seq_len, input_dim)
    
    # 目标值是序列中每个时间步的平均值加上噪声
    y = np.mean(X, axis=1) + 0.1 * np.random.randn(n_samples, input_dim)
    
    # 转换为PyTorch张量
    X = torch.FloatTensor(X)
    y = torch.FloatTensor(y)
    
    return X, y

# 训练函数
def train_model(model, train_loader, criterion, optimizer, epochs, device='cpu'):
    model.train()
    for epoch in range(epochs):
        total_loss = 0
        for X_batch, y_batch in train_loader:  # (32,10,1)  (32,1)
            X_batch, y_batch = X_batch.to(device), y_batch.to(device)
            
            # 前向传播
            outputs = model(X_batch)  # 
            loss = criterion(outputs, y_batch)
            
            # 反向传播和优化
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
            total_loss += loss.item()
        
        # 打印每个epoch的损失
        if (epoch + 1) % 10 == 0:
            print(f'Epoch [{epoch+1}/{epochs}], Loss: {total_loss/len(train_loader):.4f}')

# 评估函数
def evaluate_model(model, test_loader, device='cpu'):
    model.eval()
    predictions = []
    targets = []
    
    with torch.no_grad():
        for X_batch, y_batch in test_loader:
            X_batch, y_batch = X_batch.to(device), y_batch.to(device)
            outputs = model(X_batch)
            
            predictions.append(outputs.cpu().numpy())
            targets.append(y_batch.cpu().numpy())
    
    # 计算均方误差
    predictions = np.concatenate(predictions)
    targets = np.concatenate(targets)
    mse = np.mean((predictions - targets) ** 2)
    print(f'Test MSE: {mse:.4f}')
    
    return predictions, targets

# 主函数
def main():
    # 设备配置
    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
    
    # 模型参数
    input_dim = 1      # 输入特征维度
    hidden_dim = 32    # 隐藏层维度
    output_dim = 1     # 输出维度
    num_layers = 1     # GRU层数
    seq_len = 10       # 序列长度
    
    # 训练参数
    batch_size = 32
    epochs = 50
    learning_rate = 0.001
    
    # 生成数据
    X, y = generate_data(seq_len=seq_len) # (1000,10,1)  (1000,1)
    
    # 划分训练集和测试集
    train_size = int(0.8 * len(X))
    X_train, X_test = X[:train_size], X[train_size:]
    y_train, y_test = y[:train_size], y[train_size:]
    
    # 创建数据加载器
    train_dataset = TensorDataset(X_train, y_train) # [((10,1),1),((10,1),1)]  
    test_dataset = TensorDataset(X_test, y_test)  # 
    train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)  # (32,10,1)  (32,1)
    test_loader = DataLoader(test_dataset, batch_size=batch_size)  # (32,10,1)  (32,1)
    
    # 初始化模型
    model = GRUModel(input_dim, hidden_dim, output_dim, num_layers).to(device) # 1   32   1     1
    
    # 定义损失函数和优化器
    criterion = nn.MSELoss()
    optimizer = optim.Adam(model.parameters(), lr=learning_rate)
    
    # 训练模型
    print("开始训练模型...")
    train_model(model, train_loader, criterion, optimizer, epochs, device)
    
    # 评估模型
    print("\n评估模型...")
    predictions, targets = evaluate_model(model, test_loader, device)
    
    # 打印前10个预测结果和真实值
    print("\n示例预测结果 vs 真实值:")
    for i in range(10):
        print(f'预测: {predictions[i][0]:.4f}, 真实: {targets[i][0]:.4f}')

if __name__ == "__main__":
    main()